package main

import (
	"errors"
	"time"

	"gopkg.in/mgo.v2"
)

// 连接mongodb数据库
var (
	mongodbAddr     string = "" //mongodb数据库地址
	mongodbName     string = "" //mongodb数据名称
	mongodbUser     string = "" //mongodb用户名
	mongodbPassword string = "" //mongodb密码
)

var (
	session *mgo.Session
)

func init() {
	mongodbAddr = "127.0.0.1"
	mongodbName = "demo"
	mongodbUser = "root"
	mongodbPassword = "ming"
}

func GetMongoSession() *mgo.Session {
	if session == nil {
		var err error

		if mongodbUser == "" || mongodbPassword == "" {
			session, err = mgo.Dial(mongodbAddr)
		} else {
			dialInfo := &mgo.DialInfo{
				Addrs:     []string{mongodbAddr},
				Direct:    false,
				Timeout:   time.Second * 30,
				Database:  mongodbName,
				Source:    "admin",
				Username:  mongodbUser,
				Password:  mongodbPassword,
				PoolLimit: 4096, // Session.SetPoolLimit
			}

			session, err = mgo.DialWithInfo(dialInfo)
		}

		if err != nil {
			return nil
		}
	}

	return session.Clone()
}

func WithMongoCollection(collectionName string, s func(*mgo.Collection) error) error {
	session := GetMongoSession()
	if session == nil {
		return errors.New("获取mongodb连接失败")
	}
	defer session.Close()

	c := session.DB(mongodbName).C(collectionName)
	return s(c)
}
